"""
此文件逻辑是用于预测5s的CSV数据集，并将预测结果保存为csv文件。

"""

from keras.models import load_model
from dataset import dataset_features_labels_for_predict as dataset
import numpy as np
import os
import pandas as pd
from utility import trim

class MultiModelPredictor_V2:  
    def __init__(self, dataset_original, file_path, model_input_headers, model_dir, prediction_column_names, save_dir, train_size=0, timestep=50,scaler_path=None):  
        self.dataset_original = dataset_original
        self.file_path = file_path
        self.model_input_headers = model_input_headers
        self.model_dir = model_dir
        self.prediction_column_names = prediction_column_names
        self.save_dir = save_dir
        self.train_size = train_size
        self.timestep = timestep
        self.scaler_path = scaler_path
        # 既然模型文件名格式已经规范为Model_<prediction_name>.h5,那么就可以直接按照 prediction_names 的顺序构造对应的模型路径 
        self.model_files = [os.path.join(model_dir, f"{name[0]}_model.h5") for name in self.prediction_column_names if os.path.exists(os.path.join(model_dir, f"{name[0]}_model.h5"))]
         # 与 model_files 一一对应的尾部缓存
        self.last_tails = [None] * len(self.model_files)
        # 模型数量、预测列名数量、输入通道数量必须一致
        if not (len(self.model_files) == len(prediction_column_names) == len(model_input_headers)):
            print("模型文件数、预测列名数量、特征通道数组数量不一致，开始裁剪")
            self.prediction_column_names, self.model_input_headers=trim._align_to_model_files(self.model_files, prediction_column_names, model_input_headers)
        
    def Predict_all(self,last_tail=None):
        predictions = []
 
        for i, (model_path, input_headers, label_names) in enumerate(zip(self.model_files, self.model_input_headers, self.prediction_column_names)):
           
            # 为每个模型构建一个专属数据集（只选定特征通道）
            self.dataset = dataset.LSTMDataset_Predict(self.dataset_original, feature_names=input_headers,label_names=label_names, train_size=self.train_size, timestep=self.timestep,scaler_path=self.scaler_path)
          
            x_predict_scaled = self.dataset.get_timeclycle_data()
            # (3809, 50, 4)
            # print(x_test_scaled.shape[1])
            timestep = x_predict_scaled.shape[1]
            feature_dim = x_predict_scaled.shape[2]

            # 根据位置添加不同的前缀
            if last_tail[i] is None:
                # 第一个模型使用zero padding
                zero_padding = np.zeros((self.timestep, timestep, feature_dim))
                x_predict_scaled = np.concatenate((zero_padding, x_predict_scaled), axis=0)
            else:
                # 后续模型用上一轮的尾部作为前缀
                x_predict_scaled = np.concatenate((last_tail[i], x_predict_scaled), axis=0)

            # 更新 last_tail 为当前模型的最后 timestep 行（准备给下一个模型使用）
            self.last_tails[i] = x_predict_scaled[-self.timestep:].copy()
            model = load_model(model_path)
            # print(f"正在预测{model_path}...")
            y_pred_scaled = model.predict(x_predict_scaled)
            y_pred = self.dataset.y_scaler.inverse_transform(y_pred_scaled)
            predictions.append(y_pred.squeeze())

        return np.stack(predictions, axis=1)
        
    def save_prediction_to_csv(self, predictions):
        """
        将原始数据与预测值合并并保存到指定路径下
        """
        flat_column_names = sum(self.prediction_column_names, [])
        # print(flat_column_names)
        # print(predictions.shape)
        df_predictions = pd.DataFrame(predictions, columns=flat_column_names)
        #将两个 DataFrame 横向拼接（按列合并）成一个最终输出表
        df_output = pd.concat([self.dataset.dataset_original.reset_index(drop=True), df_predictions], axis=1)
        # --- 关键新增：根据 file_path 提取日期文件夹 ---
        date_folder = os.path.basename(os.path.dirname(self.file_path))   # e.g. "05_17"
        # 合并预测文件夹和日期文件夹路径即为最终保存的路径
        save_folder = os.path.join(self.save_dir, date_folder)
        os.makedirs(save_folder, exist_ok=True)
        # 获取原文件名
        filename = os.path.basename(self.file_path)
        # print(filename)
        save_path = os.path.join(save_folder, filename)

        df_output.to_csv(save_path, index=False, encoding='utf-8-sig')
        print(f"已保存预测结果至：{save_path}")

if __name__ == '__main__':
    # 这里第一列的列名不要写进来
    pass